package tech.waterism.model;

import cn.hutool.core.io.resource.ResourceUtil;
import com.alibaba.fastjson.JSON;
import org.apache.log4j.Logger;
import tech.waterism.modelbase.PredictResult;

import java.util.*;

/**
 * SMS_3 和 LAG_3 串联运行 Pipeline
 * SMS_3 的产流输出作为 LAG_3 的汇流输入
 * 
 * <p>工作流程：</p>
 * <ol>
 *   <li>读取 SMS_3 输入数据</li>
 *   <li>运行 SMS_3 产流模型</li>
 *   <li>将 SMS_3 的三水源产流（RSim）转换为 LAG_3 的输入格式（flowSeries with rSim）</li>
 *   <li>合并 LAG_3 参数数据</li>
 *   <li>运行 LAG_3 汇流模型</li>
 *   <li>输出最终流量过程</li>
 * </ol>
 *
 * @author AI Assistant
 * @version 1.0
 * @since 2025-01-25
 */
public class SMS3_LAG3_Pipeline {

    private static final Logger logger = Logger.getLogger(SMS3_LAG3_Pipeline.class);

    /**
     * 执行 SMS_3 + LAG_3 串联计算（使用 Map 输入）
     * 
     * @param sms3InputData SMS_3 输入数据（包含模型参数、时序数据等）
     * @return LAG_3 的模型实例（包含 predict 结果和 lastOutput）
     */
    public static LAG_3 runPipelineWithData(Map<String, Object> sms3InputData) {
        logger.info("========== 开始 SMS_3 → LAG_3 串联计算 ==========");
        
        // ===== 步骤1: 运行 SMS_3 产流模型 =====
        logger.info("步骤1: 运行 SMS_3 产流模型...");
        SMS_3 sms3Model = new SMS_3(sms3InputData);
        PredictResult sms3Result = sms3Model.predict();
        
        logger.info("SMS_3 产流计算完成:");
        logger.info("  - 时段数: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[0].length : 0));
        logger.info("  - RS（地表径流）: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[0].length + " 个值" : "null"));
        logger.info("  - RI（壤中流）: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[1].length + " 个值" : "null"));
        logger.info("  - RG（地下径流）: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[2].length + " 个值" : "null"));
        
        // ===== 步骤2: 构建 LAG_3 输入数据 =====
        logger.info("\n步骤2: 构建 LAG_3 输入数据...");
        
        // 从 SMS_3 输入中提取 LAG_3 参数（如果有的话）
        Map<String, Object> lag3Data = extractLAG3Parameters(sms3InputData);
        
        // 将 SMS_3 的输出转换为 LAG_3 的 flowSeries 格式
        List<Map<String, Object>> flowSeries = convertSMS3OutputToFlowSeries(
            sms3Result.getDt(),              // 时间序列
            sms3Result.getRain(),            // 降雨序列
            sms3Result.getRunoffSim(),       // 总产流序列
            sms3Result.getRSim()             // 三水源产流 [RS, RI, RG]
        );
        
        lag3Data.put("flowSeries", flowSeries);
        
        logger.info("LAG_3 输入数据构建完成:");
        logger.info("  - flowSeries 时段数: " + flowSeries.size());
        
        // ===== 步骤3: 运行 LAG_3 汇流模型 =====
        logger.info("\n步骤3: 运行 LAG_3 汇流模型...");
        LAG_3 lag3Model = new LAG_3(lag3Data);
        PredictResult lag3Result = lag3Model.predict();
        
        logger.info("LAG_3 汇流计算完成:");
        logger.info("  - 流量过程 QSim: " + (lag3Result.getQSim() != null ? lag3Result.getQSim().length + " 个时段" : "null"));
        
        if (lag3Result.getQSim() != null && lag3Result.getQSim().length > 0) {
            logger.info("  - 前5个流量值:");
            for (int i = 0; i < Math.min(5, lag3Result.getQSim().length); i++) {
                logger.info(String.format("    时段 %d: %.3f m³/s", i + 1, lag3Result.getQSim()[i]));
            }
            
            // 找出峰值流量
            double maxQ = 0;
            int maxQIndex = 0;
            for (int i = 0; i < lag3Result.getQSim().length; i++) {
                if (lag3Result.getQSim()[i] > maxQ) {
                    maxQ = lag3Result.getQSim()[i];
                    maxQIndex = i;
                }
            }
            logger.info(String.format("  - 峰值流量: %.3f m³/s (第 %d 时段)", maxQ, maxQIndex + 1));
        }
        
        logger.info("\n========== SMS_3 → LAG_3 串联计算完成 ==========");
        
        // 返回 LAG_3 模型实例（包含 lastOutput）
        return lag3Model;
    }
    
    /**
     * 执行 SMS_3 + LAG_3 串联计算（使用文件路径）
     * 
     * @param sms3InputPath SMS_3 输入文件路径（resources下的路径）
     * @param lag3InputPath LAG_3 参数文件路径（resources下的路径）
     * @return LAG_3 的模型实例（包含 predict 结果和 lastOutput）
     */
    @SuppressWarnings("unchecked")
    public static LAG_3 runPipeline(String sms3InputPath, String lag3InputPath) {
        logger.info("========== 开始 SMS_3 → LAG_3 串联计算 ==========");
        
        // ===== 步骤1: 运行 SMS_3 产流模型 =====
        logger.info("步骤1: 读取 SMS_3 输入并运行产流模型...");
        String sms3Input = ResourceUtil.readUtf8Str(sms3InputPath);
        Map<String, Object> sms3Data = JSON.parseObject(sms3Input, Map.class);
        
        SMS_3 sms3Model = new SMS_3(sms3Data);
        PredictResult sms3Result = sms3Model.predict();
        
        logger.info("SMS_3 产流计算完成:");
        logger.info("  - 时段数: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[0].length : 0));
        logger.info("  - RS（地表径流）: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[0].length + " 个值" : "null"));
        logger.info("  - RI（壤中流）: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[1].length + " 个值" : "null"));
        logger.info("  - RG（地下径流）: " + (sms3Result.getRSim() != null ? sms3Result.getRSim()[2].length + " 个值" : "null"));
        
        // ===== 步骤2: 构建 LAG_3 输入数据 =====
        logger.info("\n步骤2: 构建 LAG_3 输入数据...");
        
        // 读取 LAG_3 参数文件
        String lag3Input = ResourceUtil.readUtf8Str(lag3InputPath);
        Map<String, Object> lag3Data = JSON.parseObject(lag3Input, Map.class);
        
        // 将 SMS_3 的输出转换为 LAG_3 的 flowSeries 格式
        List<Map<String, Object>> flowSeries = convertSMS3OutputToFlowSeries(
            sms3Result.getDt(),              // 时间序列
            sms3Result.getRain(),            // 降雨序列
            sms3Result.getRunoffSim(),       // 总产流序列
            sms3Result.getRSim()             // 三水源产流 [RS, RI, RG]
        );
        
        lag3Data.put("flowSeries", flowSeries);
        
        logger.info("LAG_3 输入数据构建完成:");
        logger.info("  - flowSeries 时段数: " + flowSeries.size());
        
        // ===== 步骤3: 运行 LAG_3 汇流模型 =====
        logger.info("\n步骤3: 运行 LAG_3 汇流模型...");
        LAG_3 lag3Model = new LAG_3(lag3Data);
        PredictResult lag3Result = lag3Model.predict();
        
        logger.info("LAG_3 汇流计算完成:");
        logger.info("  - 流量过程 QSim: " + (lag3Result.getQSim() != null ? lag3Result.getQSim().length + " 个时段" : "null"));
        
        if (lag3Result.getQSim() != null && lag3Result.getQSim().length > 0) {
            logger.info("  - 前5个流量值:");
            for (int i = 0; i < Math.min(5, lag3Result.getQSim().length); i++) {
                logger.info(String.format("    时段 %d: %.3f m³/s", i + 1, lag3Result.getQSim()[i]));
            }
            
            // 找出峰值流量
            double maxQ = 0;
            int maxQIndex = 0;
            for (int i = 0; i < lag3Result.getQSim().length; i++) {
                if (lag3Result.getQSim()[i] > maxQ) {
                    maxQ = lag3Result.getQSim()[i];
                    maxQIndex = i;
                }
            }
            logger.info(String.format("  - 峰值流量: %.3f m³/s (第 %d 时段)", maxQ, maxQIndex + 1));
        }
        
        logger.info("\n========== SMS_3 → LAG_3 串联计算完成 ==========");
        
        // 返回 LAG_3 模型实例（包含 lastOutput）
        return lag3Model;
    }
    
    /**
     * 从 SMS_3 输入数据中提取 LAG_3 所需参数
     * 
     * @param sms3InputData SMS_3 输入数据
     * @return LAG_3 参数 Map
     */
    private static Map<String, Object> extractLAG3Parameters(Map<String, Object> sms3InputData) {
        Map<String, Object> lag3Params = new HashMap<>();
        
        // LAG_3 必需参数及其默认值
        lag3Params.put("KK", sms3InputData.getOrDefault("KK", 3.0));
        lag3Params.put("MP", sms3InputData.getOrDefault("MP", "4"));
        lag3Params.put("CG", sms3InputData.getOrDefault("CG", "0.968"));
        lag3Params.put("CI", sms3InputData.getOrDefault("CI", "0.571"));
        lag3Params.put("CS", sms3InputData.getOrDefault("CS", "0.103"));
        lag3Params.put("LAG", sms3InputData.getOrDefault("LAG", 0.0));
        lag3Params.put("X", sms3InputData.getOrDefault("X", "-0.803"));
        lag3Params.put("clen", sms3InputData.getOrDefault("clen", 3));
        
        // 流域面积（必需参数，从 parameters 对象或顶层获取）
        Object fValue = null;
        
        // 优先从 parameters 对象中获取
        if (sms3InputData.containsKey("parameters")) {
            @SuppressWarnings("unchecked")
            Map<String, Object> parameters = (Map<String, Object>) sms3InputData.get("parameters");
            if (parameters != null && parameters.containsKey("F")) {
                fValue = parameters.get("F");
            }
        }
        
        // 如果 parameters 中没有，尝试从顶层获取
        if (fValue == null && sms3InputData.containsKey("F")) {
            fValue = sms3InputData.get("F");
        }
        
        // 如果还是没有，抛出异常
        if (fValue == null) {
            throw new IllegalArgumentException("缺少必需字段: F (流域面积)，请在 parameters 对象或顶层提供");
        }
        
        lag3Params.put("F", fValue);
        
        // 初始状态参数
        lag3Params.put("QSP", sms3InputData.getOrDefault("QSP", "0"));
        lag3Params.put("QIP", sms3InputData.getOrDefault("QIP", "0"));
        lag3Params.put("QGP", sms3InputData.getOrDefault("QGP", "0"));
        
        // 初始河道流量序列
        if (sms3InputData.containsKey("QSIG")) {
            lag3Params.put("QSIG", sms3InputData.get("QSIG"));
        } else {
            lag3Params.put("QSIG", Arrays.asList(0, 0, 0, 0, 0, 0));
        }
        
        // 初始河段流量
        if (sms3InputData.containsKey("QXSIG")) {
            lag3Params.put("QXSIG", sms3InputData.get("QXSIG"));
        } else {
            lag3Params.put("QXSIG", Arrays.asList(0));
        }
        
        // 起始时间
        if (sms3InputData.containsKey("start")) {
            lag3Params.put("start", sms3InputData.get("start"));
        }
        
        logger.info("LAG_3 参数提取完成:");
        logger.info("  - KK: " + lag3Params.get("KK"));
        logger.info("  - CS: " + lag3Params.get("CS"));
        logger.info("  - CI: " + lag3Params.get("CI"));
        logger.info("  - CG: " + lag3Params.get("CG"));
        logger.info("  - F: " + lag3Params.get("F"));
        
        return lag3Params;
    }
    
    /**
     * 将 SMS_3 的输出转换为 LAG_3 的 flowSeries 输入格式
     * 
     * <p>flowSeries 格式：</p>
     * <pre>
     * {
     *   "time": "2025-06-05 08:00:00",
     *   "rain": 10.5,
     *   "runoff": 5.2,
     *   "rSim": [RS, RI, RG]
     * }
     * </pre>
     * 
     * @param timeArray 时间序列
     * @param rainArray 降雨序列 (mm)
     * @param runoffArray 总产流序列 (mm)
     * @param rSimArray 三水源产流 [RS, RI, RG] (mm)
     * @return flowSeries 列表
     */
    public static List<Map<String, Object>> convertSMS3OutputToFlowSeries(
            Date[] timeArray,
            double[] rainArray,
            double[] runoffArray,
            double[][] rSimArray) {
        
        logger.info("开始转换 SMS_3 输出为 LAG_3 flowSeries 格式...");
        
        List<Map<String, Object>> flowSeries = new ArrayList<>();
        java.text.SimpleDateFormat sdf = new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
        
        int length = timeArray != null ? timeArray.length : 0;
        
        for (int i = 0; i < length; i++) {
            Map<String, Object> point = new HashMap<>();
            
            // 时间
            point.put("time", sdf.format(timeArray[i]));
            
            // 降雨
            if (rainArray != null && i < rainArray.length) {
                point.put("rain", rainArray[i]);
            } else {
                point.put("rain", 0.0);
            }
            
            // 总产流
            if (runoffArray != null && i < runoffArray.length) {
                point.put("runoff", runoffArray[i]);
            } else {
                point.put("runoff", 0.0);
            }
            
            // 三水源产流 [RS, RI, RG]
            if (rSimArray != null && rSimArray.length >= 3) {
                List<Double> rSim = new ArrayList<>();
                rSim.add(i < rSimArray[0].length ? rSimArray[0][i] : 0.0);  // RS
                rSim.add(i < rSimArray[1].length ? rSimArray[1][i] : 0.0);  // RI
                rSim.add(i < rSimArray[2].length ? rSimArray[2][i] : 0.0);  // RG
                point.put("rSim", rSim);
            } else {
                point.put("rSim", Arrays.asList(0.0, 0.0, 0.0));
            }
            
            flowSeries.add(point);
        }
        
        logger.info("转换完成，共 " + flowSeries.size() + " 个时段");
        return flowSeries;
    }
    
    /**
     * 测试入口 - 使用标准JSON文件
     */
    public static void main(String[] args) {
        logger.info("========================================");
        logger.info("SMS_3 + LAG_3 串联运行测试");
        logger.info("========================================\n");
        
        try {
            // 执行串联计算
            LAG_3 lag3Model = runPipeline(
                "json/SMS_3-input.json",  // SMS_3 输入文件
                "json/LAG-input.json"     // LAG_3 参数文件
            );
            
            // 获取标准化输出对象
            tech.waterism.model.dto.ModelOutputSimple output = lag3Model.getLastOutput();
            
            // 输出结果摘要
            logger.info("\n========== 最终结果摘要 ==========");
            logger.info("模型名称: " + output.getModelName());
            logger.info("计算完成时间: " + output.getComputeTime());
            if (output.getTimeSeries() != null && !output.getTimeSeries().isEmpty()) {
                logger.info("时间序列长度: " + output.getTimeSeries().size() + " 个时段");
                logger.info("前5个时段的流量:");
                for (int i = 0; i < Math.min(5, output.getTimeSeries().size()); i++) {
                    Map<String, Object> point = output.getTimeSeries().get(i);
                    logger.info(String.format("  时段 %d: time=%s, discharge=%.3f m³/s", 
                        i + 1, point.get("time"), point.get("discharge")));
                }
            }
            
            // 输出完整JSON（用于调试）
            logger.info("\n完整输出JSON:");
            System.out.println(JSON.toJSONString(output));
            
        } catch (Exception e) {
            logger.error("串联计算过程中发生错误", e);
            e.printStackTrace();
        }
    }
}

